# python -m lilab.OpenLabCluster_train.a1_mirror_mutual_merge_clippredpkl A.clippredpkl
import numpy as np
import pickle
import os
import os.path as osp
import pandas as pd
import argparse

from lilab.openlabcluster_postprocess.s1_merge_3_file import get_assert_1_file
from lilab.openlabcluster_postprocess.s1a_clipNames_inplace_parse import parse_name
from lilab.OpenLabCluster_train.a1_mirror_mutual_filt_clippredpkl import factory_label_mirror_equal_start0

# clippredpkl = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32_k36/semiseq2seq_iter0/output/far_ns_with_s_recluster_k36_refine/representitive_k36_filt_perc91/Representive_K36.clippredpkl'

def main(clippredpkl):
    project_dir = osp.dirname(clippredpkl)
    clippreddata = pickle.load(open(clippredpkl, 'rb'))
    if 'df_clipNames' not in clippreddata:
        clippreddata['df_clipNames'] = parse_name(clippreddata['clipNames'])
    df_clipNames = clippreddata['df_clipNames']

    nK_mutual = clippreddata['nK_mutual']
    nK_mirror_half = clippreddata.get('nK_mirror_half', clippreddata.get('nK_mirrorhalf',None))
    cluster_labels = clippreddata['cluster_labels']
    assert cluster_labels.max() == nK_mutual + nK_mirror_half*2

    fun_label_equal = factory_label_mirror_equal_start0(nK_mutual, nK_mirror_half, start=1)
    cluster_labels_part = fun_label_equal(cluster_labels)
    #%% save data
    clippreddataNew = dict()
    cluster_labels_part_1start = cluster_labels_part - cluster_labels_part.min() + 1 #start from 1
    nK = cluster_labels_part_1start.max()
    assert nK == nK_mutual + nK_mirror_half
    clippreddataNew['ncluster'] = nK
    clippreddataNew['ntwin'] = clippreddata['ntwin']
    clippreddataNew['cluster_labels'] = cluster_labels_part_1start
    clippreddataNew['cluster_names'] = clippreddata['cluster_names_mutualmerge']
    clippreddataNew['embedding'] = clippreddata['embedding']*np.nan
    clippreddataNew['embedding_d2'] = clippreddata['embedding_d2']*np.nan
    clippreddataNew['clipNames'] = clippreddata['clipNames']
    clippreddataNew['df_clipNames'] = df_clipNames
    clippreddataNew['nK_mutual'] = nK_mutual
    clippreddataNew['nK_mirror_half'] = nK_mirror_half

    output_dir = osp.join(project_dir, f'mirror_merge')
    os.makedirs(output_dir, exist_ok=True)
    clippredpklNew = osp.join(output_dir, f'mirror_merge_K{nK}.clippredpkl')
    print(f'save to {clippredpklNew}')
    pickle.dump(clippreddataNew, open(clippredpklNew, 'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('clippredpkl', type=str, help='clippredpkl file')
    args = parser.parse_args()

    main(args.clippredpkl)
    
